from data_preprocess_jl import have_test_preprocess_data, test_data_analysis_encode_save, load_test_preprocess_data
import joblib
from sklearn.metrics import roc_auc_score, classification_report
from sklearn.metrics import confusion_matrix, roc_curve, RocCurveDisplay, roc_auc_score
import matplotlib.pyplot as plt
from utils.log_jl import Logger
from datetime import datetime
import numpy as np


class Predict_XGBoost_TL(object):
    def __init__(self):
        if have_test_preprocess_data() is False:
            test_data_analysis_encode_save()
        self.datasource = load_test_preprocess_data()
        # 配置日志记录
        logfile_name = 'predict_' + str(datetime.now().strftime('%Y%m%d%H%M%S'))
        self.logfile = Logger('../', logfile_name).get_logger()
        self.original_model = None
        self.reduced_model = None

    def predict_xgb_model_from_localdata(self):
        x = self.datasource.iloc[:, 1:]
        y = self.datasource.iloc[:, 0]
        self.original_model = joblib.load('../output/model/xgboost_model_jl.pkl')
        y_pred = self.original_model.predict(x)
        print(f'{classification_report(y, y_pred)}')
        self.logfile.info(f'{classification_report(y, y_pred)}')
        y_pred_proba = self.original_model.predict_proba(x)[:, 1]
        print(f"ROC AUC: {roc_auc_score(y, y_pred_proba)}")
        self.logfile.info(f"ROC AUC: {roc_auc_score(y, y_pred_proba)}")
        # 混淆矩阵查看类别分布
        cm = confusion_matrix(y, y_pred)
        print("Confusion Matrix:\n", cm)

    def plot_roc_curves(self):
        x = self.datasource.iloc[:, 1:]
        y = self.datasource.iloc[:, 0]

        # 获取预测概率
        y_pred_proba_original = self.get_y_pred_proba(self.original_model, x)
        y_pred_proba_reduced = self.get_y_pred_proba(self.reduced_model, x)

        # 计算FPR  和 TPR
        fpr_original, tpr_original, thresholds_original = roc_curve(y, y_pred_proba_original)
        fpr_reduced, tpr_reduced, thresholds_reduced = roc_curve(y, y_pred_proba_reduced)

        # 计算AUC
        auc_original = roc_auc_score(y, y_pred_proba_original)
        auc_reduced = roc_auc_score(y, y_pred_proba_reduced)

        # 绘图
        plt.figure(figsize=(8, 6))
        plt.plot(fpr_original, tpr_original, label=f'原始模型 AUC = {auc_original:.4f}')
        plt.plot(fpr_reduced, tpr_reduced, label=f'精简模型 AUC = {auc_reduced:.4f}')
        plt.plot([0, 1], [0, 1], 'k--', label='随机猜测')
        plt.xlabel('False Positive Rate (FPR)')
        plt.ylabel('True Positive Rate (TPR)')
        plt.title('ROC 曲线对比')
        plt.legend(loc='lower right')
        plt.grid(True)
        plt.tight_layout()
        plt.savefig('../output/figure/原模型和精简模型ROC曲线对比图.png')
        plt.show()

        # 增加输出report和混淆矩阵
        y_pred_original = self.get_y_pred(self.original_model, x)
        y_pred_reduced = self.get_y_pred(self.reduced_model, x)
        print(f'原始模型报告：\n{classification_report(y, y_pred_original)}')
        print(f'原始模型混淆矩阵：\n{confusion_matrix(y, y_pred_original)}')
        print(f'精简模型报告：\n{classification_report(y, y_pred_reduced)}')
        print(f'精简模型混淆矩阵：\n{confusion_matrix(y, y_pred_reduced)}')

    def roc_graph_single(self, model=None, name=None):
        '''
        绘制指定模型的ROC曲线
        :param model: 需要绘制ROC曲线的模型
        :return:
        '''
        x = self.datasource.iloc[:, 1:]
        y = self.datasource.iloc[:, 0]
        common_features = [col for col in x.columns if col in model.feature_names_in_]
        # 绘制 ROC 曲线
        RocCurveDisplay.from_estimator(model, x[common_features], y, name=name, plot_chance_level=True)
        # 保存到本地
        plt.savefig(f'../output/figure/{name}的ROC曲线图.png')
        plt.show()

    def roc_graph_original_model(self):
        if self.original_model is None or self.reduced_model is None:
            self.logfile.info("Please train the model first.")
            print("请先执行加载模型的方法")
            return
        self.roc_graph_single(self.original_model, name='原始模型')

    def roc_graph_reduced_model(self):
        if self.reduced_model is None:
            self.logfile.info("Please train the model first.")
            print("请先执行加载模型的方法")
            return
        self.roc_graph_single(self.reduced_model, name='精简模型')

    def get_y_pred_proba(self, model, x):
        try:
            y_pred_proba = model.predict_proba(x)[:, 1]
            return y_pred_proba
        except ValueError as e:
            # 如果特征不匹配，说明需要筛选特征
            common_features = [col for col in x.columns if col in model.feature_names_in_]
            y_pred_proba = model.predict_proba(x[common_features])[:, 1]
            return y_pred_proba

    def get_y_pred(self, model, x):
        try:
            y_pred = model.predict(x)
            return y_pred
        except ValueError as e:
            # 如果特征不匹配，说明需要筛选特征
            common_features = [col for col in x.columns if col in model.feature_names_in_]
            y_pred = model.predict(x[common_features])
            return y_pred

    def compare_model_auc(self, original_model_path, reduced_model_path):
        """
        比较两个模型在测试集上的 AUC 表现
        :param original_model_path: 原始模型路径
        :param reduced_model_path: 精简模型路径
        """
        x = self.datasource.iloc[:, 1:]
        y = self.datasource.iloc[:, 0]

        # 加载模型
        self.original_model = joblib.load(original_model_path)
        self.reduced_model = joblib.load(reduced_model_path)

        # 原始模型预测
        y_pred_proba_original = self.get_y_pred_proba(self.original_model, x)

        # 精简模型预测
        y_pred_proba_reduced = self.get_y_pred_proba(self.reduced_model, x)

        auc_original = roc_auc_score(y, y_pred_proba_original)
        auc_reduced = roc_auc_score(y, y_pred_proba_reduced)

        self.logfile.info(f'原模型的AUC: {auc_original:.4f}')
        self.logfile.info(f'精简模型的AUC: {auc_reduced:.4f}')
        self.logfile.info(f'两者AUC差值: {abs(auc_reduced - auc_original):.4f}')

        print(f'原模型的AUC: {auc_original:.4f}')
        print(f'精简模型的AUC: {auc_reduced:.4f}')
        print(f'两者AUC差值: {abs(auc_reduced - auc_original):.4f}')

        return auc_original, auc_reduced
